
#include "point-cloud-v2.h"

namespace point_cloud_v2
{
	class PointCloudVars : public SessionSampleVars
	{
	public:
		SPADAS_VARS(PointCloud, SessionSampleVars)
		PointCloudData data;

		SessionMatrixSample toMatrixSample()
		{
			UInt iCol = UINF, rCol = UINF, gCol = UINF, bCol = UINF, lidCol = UINF, toCol = UINF, yaCol = UINF, midCol = UINF;
			Array<UInt> byteFieldColTable(100, UINF); // 100~199
			Array<UInt> shortFieldColTable(100, UINF); // 200~299
			Array<UInt> floatFieldColTable(100, UINF); // 300~399

			UInt nPoints = data.pointCount;
			if (data.points.nDims() != 2 || data.points.dimAt(1) != 3) nPoints = 0;
			else if (nPoints > data.points.dimAt(0)) nPoints = data.points.dimAt(0);

			UInt nCols = 3;
			if (nPoints > 0)
			{
				if (data.intensities.size() == nPoints) iCol = nCols++;
				if (data.colors.size() == nPoints)
				{
					rCol = nCols++;
					gCol = nCols++;
					bCol = nCols++;
				}
				if (data.laserIDs.size() == nPoints) lidCol = nCols++;
				if (data.timeOffsets.size() == nPoints) toCol = nCols++;
				if (data.yawAngles.size() == nPoints) yaCol = nCols++;
				if (data.mirrorIDs.size() == nPoints) midCol = nCols++;
				for (auto pair = data.byteValues.keyValues().firstElem(); pair.valid(); ++pair)
				{
					Int fieldID = pair->key.value();
					if (fieldID < 100 || fieldID >= 200 || pair->value.size() != nPoints) continue;
					byteFieldColTable[fieldID - 100] = nCols++;
				}
				for (auto pair = data.shortValues.keyValues().firstElem(); pair.valid(); ++pair)
				{
					Int fieldID = pair->key.value();
					if (fieldID < 200 || fieldID >= 300 || pair->value.size() != nPoints) continue;
					shortFieldColTable[fieldID - 200] = nCols++;
				}
				for (auto pair = data.floatValues.keyValues().firstElem(); pair.valid(); ++pair)
				{
					Int fieldID = pair->key.value();
					if (fieldID < 300 || fieldID >= 400 || pair->value.size() != nPoints) continue;
					floatFieldColTable[fieldID - 300] = nCols++;
				}
			}

			SessionMatrixSample sample(Size2D::dims(nPoints + 1, nCols));
			sample.timestamp() = timestamp;

			Float *headData = sample.matData().data();
			Float *pointData = &sample.matData()[nCols];
			
			Float *xyzData = data.points.data();
			UInt k = 0;
			for (UInt a = 0; a < nPoints; a++)
			{
				pointData[a * nCols] = xyzData[k++];
				pointData[a * nCols + 1] = xyzData[k++];
				pointData[a * nCols + 2] = xyzData[k++];
			}

			if (iCol != UINF)
			{
				headData[iCol] = 1;
				Float *iData = data.intensities.data();
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols + iCol] = iData[a];
				}
			}

			if (rCol != UINF)
			{
				headData[rCol] = 2;
				headData[gCol] = 3;
				headData[bCol] = 4;
				ColorRGB *colorData = data.colors.data();
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols + rCol] = colorData[a].r;
					pointData[a * nCols + gCol] = colorData[a].g;
					pointData[a * nCols + bCol] = colorData[a].b;
				}
			}

			if (lidCol != UINF)
			{
				headData[lidCol] = 5;
				Word *lidData = data.laserIDs.data();
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols + lidCol] = lidData[a];
				}
			}

			if (toCol != UINF)
			{
				headData[toCol] = 6;
				Float *toData = data.timeOffsets.data();
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols + toCol] = toData[a];
				}
			}

			if (yaCol != UINF)
			{
				headData[yaCol] = 7;
				Float *yaData = data.yawAngles.data();
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols + yaCol] = yaData[a];
				}
			}

			if (midCol != UINF)
			{
				headData[midCol] = 8;
				Byte *midData = data.mirrorIDs.data();
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols + midCol] = midData[a];
				}
			}

			for (UInt i = 0; i < byteFieldColTable.size(); i++)
			{
				if (byteFieldColTable[i] == UINF) continue;
				UInt byteFieldCol = byteFieldColTable[i];
				headData[byteFieldCol] = 100 + i;
				Byte *byteFieldData = data.byteValues[IntKey(100 + i)].data();
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols + byteFieldCol] = byteFieldData[a];
				}
			}

			for (UInt i = 0; i < shortFieldColTable.size(); i++)
			{
				if (shortFieldColTable[i] == UINF) continue;
				UInt shortFieldCol = shortFieldColTable[i];
				headData[shortFieldCol] = 200 + i;
				Short *shortFieldData = data.shortValues[IntKey(200 + i)].data();
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols + shortFieldCol] = shortFieldData[a];
				}
			}

			for (UInt i = 0; i < floatFieldColTable.size(); i++)
			{
				if (floatFieldColTable[i] == UINF) continue;
				UInt floatFieldCol = floatFieldColTable[i];
				headData[floatFieldCol] = 300 + i;
				Float *floatFieldData = data.floatValues[IntKey(300 + i)].data();
				for (UInt a = 0; a < nPoints; a++)
				{
					pointData[a * nCols + floatFieldCol] = floatFieldData[a];
				}
			}

			return sample;
		}

		Bool fromSample(String protocol, SessionSample sample) override
		{
			if (!sample.is<SessionMatrixSample>()) return FALSE;

			auto matrixSample = sample.as<SessionMatrixSample>();
			timestamp = matrixSample.timestamp();
			const UInt nRows = matrixSample.rows();
			const UInt nCols = matrixSample.cols();
			if (nCols < 3) return FALSE;

			Float *pointData = NULL;
			UInt nPoints = 0;
			UInt iCol = UINF, rCol = UINF, gCol = UINF, bCol = UINF, lidCol = UINF, toCol = UINF, yaCol = UINF, midCol = UINF;
			Bool hasColor = FALSE;
			Array<UInt> byteFieldColTable(100, UINF); // 100~199
			Array<UInt> shortFieldColTable(100, UINF); // 200~299
			Array<UInt> floatFieldColTable(100, UINF); // 300~399

			if (protocol == "point-cloud-v2")
			{
				if (nRows == 0) return FALSE;

				pointData = matrixSample.matData().data();
				nPoints = nRows - 1;

				for (UInt a = 3; a < nCols; a++)
				{
					UInt fieldType = (UInt)pointData[a];
					switch (fieldType)
					{
					case 1:
						iCol = a;
						break;
					case 2:
						rCol = a;
						break;
					case 3:
						gCol = a;
						break;
					case 4:
						bCol = a;
						break;
					case 5:
						lidCol = a;
						break;
					case 6:
						toCol = a;
						break;
					case 7:
						yaCol = a;
						break;
					case 8:
						midCol = a;
						break;
					default:
						if (fieldType >= 100 && fieldType < 200)
						{
							byteFieldColTable[fieldType - 100] = a;
						}
						else if (fieldType >= 200 && fieldType < 300)
						{
							shortFieldColTable[fieldType - 200] = a;
						}
						else if (fieldType >= 300 && fieldType < 400)
						{
							floatFieldColTable[fieldType - 300] = a;
						}
						break;
					}
				}

				pointData = &pointData[nCols];
				hasColor = rCol != UINF && gCol != UINF && bCol != UINF;
			}
			else if (protocol == "point-cloud-v1")
			{
				pointData = matrixSample.matData().data();
				nPoints = nRows;
			}
			else return FALSE;

			data.pointCount = nPoints;
			data.points = nPoints > 0 ? FloatMat(nPoints, 3) : FloatMat();
			data.intensities = Array<Float>();
			data.colors = Array<ColorRGB>();
			data.laserIDs = Array<Word>();
			data.timeOffsets = Array<Float>();
			data.yawAngles = Array<Float>();
			data.mirrorIDs = Array<Byte>();
			data.byteValues = Map<IntKey, Array<Byte> >();
			data.shortValues = Map<IntKey, Array<Short> >();
			data.floatValues = Map<IntKey, Array<Float> >();

			if (nPoints > 0)
			{
				Float *xyzData = data.points.data();
				UInt k = 0;
				for (UInt a = 0; a < nPoints; a++)
				{
					xyzData[k++] = pointData[a * nCols];
					xyzData[k++] = pointData[a * nCols + 1];
					xyzData[k++] = pointData[a * nCols + 2];
				}

				if (iCol != UINF)
				{
					data.intensities = Array<Float>(nPoints);
					Float *iData = data.intensities.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						iData[a] = pointData[a * nCols + iCol];
					}
				}

				if (hasColor)
				{
					data.colors = Array<ColorRGB>(nPoints);
					ColorRGB *colorData = data.colors.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						colorData[a].r = (Byte)pointData[a * nCols + rCol];
						colorData[a].g = (Byte)pointData[a * nCols + gCol];
						colorData[a].b = (Byte)pointData[a * nCols + bCol];
					}
				}

				if (lidCol != UINF)
				{
					data.laserIDs = Array<Word>(nPoints);
					Word *lidData = data.laserIDs.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						lidData[a] = (Word)pointData[a * nCols + lidCol];
					}
				}

				if (toCol != UINF)
				{
					data.timeOffsets = Array<Float>(nPoints);
					Float *toData = data.timeOffsets.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						toData[a] = pointData[a * nCols + toCol];
					}
				}

				if (yaCol != UINF)
				{
					data.yawAngles = Array<Float>(nPoints);
					Float *yaData = data.yawAngles.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						yaData[a] = pointData[a * nCols + yaCol];
					}
				}

				if (midCol != UINF)
				{
					data.mirrorIDs = Array<Byte>(nPoints);
					Byte *midData = data.mirrorIDs.data();
					for (UInt a = 0; a < nPoints; a++)
					{
						midData[a] = (Byte)pointData[a * nCols + midCol];
					}
				}

				for (UInt x = 0; x < 100; x++)
				{
					if (byteFieldColTable[x] != UINF)
					{
						UInt byteFieldCol = byteFieldColTable[x];
						Array<Byte> byteValues(nPoints);
						Byte *byteFieldData = byteValues.data();
						for (UInt a = 0; a < nPoints; a++)
						{
							byteFieldData[a] = (Byte)pointData[a * nCols + byteFieldCol];
						}
						data.byteValues[IntKey(100 + x)] = byteValues;
					}
				}

				for (UInt x = 0; x < 100; x++)
				{
					if (shortFieldColTable[x] != UINF)
					{
						UInt shortFieldCol = shortFieldColTable[x];
						Array<Short> shortValues(nPoints);
						Short *shortFieldData = shortValues.data();
						for (UInt a = 0; a < nPoints; a++)
						{
							shortFieldData[a] = (Short)pointData[a * nCols + shortFieldCol];
						}
						data.shortValues[IntKey(200 + x)] = shortValues;
					}
				}

				for (UInt x = 0; x < 100; x++)
				{
					if (floatFieldColTable[x] != UINF)
					{
						UInt floatFieldCol = floatFieldColTable[x];
						Array<Float> floatValues(nPoints);
						Float *floatFieldData = floatValues.data();
						for (UInt a = 0; a < nPoints; a++)
						{
							floatFieldData[a] = pointData[a * nCols + floatFieldCol];
						}
						data.floatValues[IntKey(300 + x)] = floatValues;
					}
				}
			}

			return TRUE;
		}
	};
}

using namespace point_cloud_v2;

PointCloud::PointCloud() : Object<PointCloudVars>(new PointCloudVars, TRUE)
{}

FullTimestamp& PointCloud::timestamp()
{
	return vars->timestamp;
}

PointCloudData *PointCloud::operator ->()
{
    return &vars->data;
}

String PointCloud::getProtocolName(UInt channel)
{
    return (String)"point-cloud-v2@" + channel;
}

SessionMatrixSample PointCloud::toMatrixSample()
{
    return vars->toMatrixSample();
}

Array<String> PointCloud::getProtocolNames(UInt channel)
{
	return Array<String>::create(2, "point-cloud-v1", "point-cloud-v2").convert<String>([channel](auto& s){ return s + "@" + channel; });
}

Bool PointCloud::fromMatrixSample(String protocol, SessionMatrixSample matrixSample)
{
    return vars->fromSample(protocol, matrixSample.as<SessionSample>());
}